import argparse
import json
import logging
import os
import subprocess
import sys

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def check_ssh_access(ip):
    try:
        if sys.version_info >= (3, 5):
            subprocess.run(['ssh', '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=5', ip, 'echo', 'SSH_OK'], check=True,
                           stdout=subprocess.DEVNULL)
        else:
            subprocess.check_call(['ssh', '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=5', ip, 'echo', 'SSH_OK'],
                                  stdout=subprocess.DEVNULL)
            # process = subprocess.Popen(['ssh', '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=5', ip, 'echo', 'SSH_OK'], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
            # process.communicate()
            # if process.returncode != 0:
            #     raise subprocess.CalledProcessError(process.returncode, process.args)
        return True
    except subprocess.CalledProcessError as e:
        logging.error("SSH connection to %s failed: %s", ip, e)
        return False


def get_device_ips(ip, master_ip, a3_super_nodes):
    device_ips = []
    command_prefix = ['hccn_tool', '-i', '-vnic' if a3_super_nodes else '-ip', '-g']
    try:
        for i in range(16 if a3_super_nodes else 8):
            command_prefix.insert(2, str(i))
            if ip != master_ip:
                command_prefix.insert(0, 'ssh')
                command_prefix.insert(1, ip)
            if sys.version_info >= (3, 5):
                result = subprocess.run(command_prefix, stdout=subprocess.PIPE, check=True)
                stdout = result.stdout
            else:
                stdout = subprocess.check_output(command_prefix, stderr=subprocess.STDOUT)
                # process = subprocess.Popen(command_prefix, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                # stdout, stderr = process.communicate()
                # if process.returncode != 0:
                #     raise subprocess.CalledProcessError(process.returncode, command_prefix, stderr)
            output = stdout.decode('utf-8').strip()
            if a3_super_nodes:
                ipaddr = [line.split(':')[1].strip() for line in output.split('\n') if 'vnic ipaddr' in line][0]
            else:
                ipaddr = output.split('\n')[0].split(':')[1].strip()
            device_ips.append(ipaddr)
            if ip != master_ip:
                command_prefix.pop(0)
                command_prefix.pop(0)
            command_prefix.pop(2)
        return device_ips
    except subprocess.CalledProcessError as e:
        logging.error("Failed to get device IPs from %s: %s", ip, e.output.decode('utf-8'))
        return []


def get_super_device_ids_and_pod_ids(ip, master_ip):
    super_device_ids = []
    super_pod_ids = []
    command_prefix = ['npu-smi', 'info', '-t', 'spod-info', '-i', '*', '-c', '*']
    try:
        for i in range(8):
            for j in range(2):
                command_prefix[5] = str(i)
                command_prefix[7] = str(j)
                if ip != master_ip:
                    command_prefix.insert(0, 'ssh')
                    command_prefix.insert(1, ip)
                if sys.version_info >= (3, 5):
                    result = subprocess.run(command_prefix, stdout=subprocess.PIPE, check=True)
                    stdout = result.stdout
                else:
                    stdout = subprocess.check_output(command_prefix, stderr=subprocess.STDOUT)
                output = stdout.decode('utf-8').strip()
                sdid = [line.split(':')[1].strip() for line in output.split('\n') if 'SDID' in line][0]
                spod_id = [line.split(':')[1].strip() for line in output.split('\n') if 'Super Pod ID' in line][0]
                super_device_ids.append(sdid)
                super_pod_ids.append(spod_id)
                if ip != master_ip:
                    command_prefix.pop(0)
                    command_prefix.pop(0)
        return super_device_ids, super_pod_ids
    except subprocess.CalledProcessError as e:
        logging.error("Failed to get super device IDs and pod IDs from %s: %s", ip, e.output.decode('utf-8'))
        return [], []


def setup_ssh_keys(server_ips, master_ip, local_public_key_path):
    for ip in server_ips:
        if ip != master_ip and not check_ssh_access(ip):
            try:
                command = [
                    "ssh-copy-id",
                    "-o", "StrictHostKeyChecking=no",
                    "-o", "UserKnownHostsFile=/dev/null",
                    "-i", local_public_key_path,
                    "root@%s" % ip
                ]
                if sys.version_info >= (3, 5):
                    subprocess.run(command, check=True)
                else:
                    subprocess.check_output(command, stderr=subprocess.STDOUT)
                if not check_ssh_access(ip):
                    logging.error("%s Failed to configure password protection", ip)
                    return False
            except subprocess.CalledProcessError as e:
                logging.error("Failed to copy SSH key to %s: %s", ip, e.output.decode('utf-8'))
                return False
    return True


def update_json(json_file, master_ip, server_id, container_ip, server_count, server_index, all_servers, a3_super_nodes):
    if os.path.exists(json_file):
        with open(json_file, 'r') as f:
            data = json.load(f)
    else:
        if a3_super_nodes:
            new_supod = {
                "super_pod_id": '',
                "server_list": []
            }
            data = {
                "version": "1.2",
                "server_count": str(server_count),
                "super_pod_list": [new_supod],
                "server_list": [],
                "status": "completed"
            }
        else:
            data = {
                "version": "1.0",
                "server_count": str(server_count),
                "server_list": [],
                "status": "completed"
            }

    if data.get("server_count") == str(server_count):
        pass
    else:
        data["server_count"] = str(server_count)

    rank_id_start = (server_index - 1) * 16 if a3_super_nodes else (server_index - 1) * 8
    server_found = False
    devices = [{"device_id": str(i), "device_ip": ip, "rank_id": str(rank_id_start + i)} for i, ip in
               enumerate(get_device_ips(server_id, master_ip, a3_super_nodes))]
    if a3_super_nodes:
        super_device_ids, super_pod_ids = get_super_device_ids_and_pod_ids(server_id, master_ip)
        for idx, dev in enumerate(devices):
            dev["super_device_id"] = super_device_ids[idx]

        server_found = False
        server_entry = {"server_id": server_id}
        for pod_list in data["super_pod_list"]:
            if any(s["server_id"] == server_id for s in pod_list["server_list"]):
                break
            else:
                pod_list["server_list"].append(server_entry)
            if not pod_list["super_pod_id"]:
                pod_list["super_pod_id"] = 'supod_%s' % super_pod_ids[0]

        for server in data["server_list"]:
            if server["server_id"] == server_id:
                server_found = True
                if container_ip:
                    server["container_ip"] = container_ip
                server["device"] = devices
                break

        if not server_found:
            new_server = {
                "server_id": server_id,
                "container_ip": container_ip if container_ip else server_id,
                "host_ip": server_id,
                "device": devices,
                "host_nic_ip": "reserve"
            }
            data["server_list"].append(new_server)
    else:
        for server in data["server_list"]:
            if server["server_id"] == server_id:
                server_found = True
                if container_ip:
                    server["container_ip"] = container_ip
                server["device"] = devices
                break

        if not server_found:
            new_server = {
                "server_id": server_id,
                "container_ip": container_ip if container_ip else server_id,
                "device": devices,
                "host_nic_ip": "reserve"
            }
            data["server_list"].append(new_server)

    data["server_list"] = sorted(data["server_list"], key=lambda x: int(x["device"][0]["rank_id"]))

    with open(json_file, 'w') as f:
        json.dump(data, f, indent=4)
    if server_index == server_count:
        os.chmod(json_file, 0o640)
        for ip, _ in all_servers:
            if ip != master_ip:
                try:
                    if sys.version_info >= (3, 5):
                        subprocess.run(["scp", json_file, f"root@%s:%s" % (ip, json_file)])
                    else:
                        subprocess.check_output(["scp", json_file, "root@%s:%s" % (ip, json_file)],
                                                stderr=subprocess.STDOUT)
                except subprocess.CalledProcessError as e:
                    logging.error("Failed to distribute file to %s: %s", ip, e.output.decode('utf-8'))


def main():
    parser = argparse.ArgumentParser(description="Update JSON file with server information.")
    parser.add_argument('--json_file', type=str, default='/root/rank_table.json')
    parser.add_argument('--all_config_file', type=str, default='./all_config.yaml')

    args = parser.parse_args()

    try:
        import yaml
        with open(args.all_config_file, 'r') as f:
            config = yaml.safe_load(f)
            all_servers = config['all_ip']
            a3_super_nodes = config.get('a3_super_nodes', False)

    except ImportError as e:
        logging.error(
            "Failed to import yaml: %s\n Please manually modify ssh_rank_table.py to assign values to all_servers and a3_super_nodes.",
            e)
        all_servers = []
        while True:
            line = input("Enter IP (192.168.x.x,192.168.x.x), press Enter to finish: ")
            if not line:
                a3_super_nodes = input("Is this a3_super_nodes? (True/False): ")
                break
            server_ip, container_ip = line.split(",")
            all_servers.append([server_ip, container_ip])
        logging.info("all_servers: %s", all_servers)

    local_private_key_path = '/root/.ssh/id_rsa'
    local_public_key_path = '/root/.ssh/id_rsa.pub'

    if not (os.path.exists(local_private_key_path) and os.path.exists(local_public_key_path)):
        try:
            subprocess.run(["ssh-keygen", "-t", "rsa"])
        except subprocess.CalledProcessError as e:
            logging.error("Failed to generate SSH keys: %s", e.output.decode('utf-8'))

    master_ip = all_servers[0][0]
    beta_pass = setup_ssh_keys([ip for ip, _ in all_servers], master_ip, local_public_key_path)
    if not beta_pass:
        return logging.info("Program terminated, please configure password protection manually.")

    for index, (server_ip, container_ip) in enumerate(all_servers, start=1):
        try:
            update_json(args.json_file, master_ip, server_ip, container_ip, len(all_servers), index, all_servers,
                        a3_super_nodes)
        except Exception as e:
            logging.error("Failed to update JSON for server %s: %s", server_ip, e)


if __name__ == "__main__":
    main()
